import torch
import torch.nn as nn
import torch.nn.functional as F


class EncoderRNN(nn.Module):
    def __init__(self, input_size, hidden_size, dropout=0.1, bidirectional=False):
        super().__init__()

        self.gru = nn.GRU(input_size, hidden_size, batch_first=True, bidirectional=bidirectional)
        self.dropout = nn.Dropout(dropout)

    def forward(self, embedded):
        """
        :param embedded: embedded of sequences
        :return:
        """
        d = self.dropout(embedded)
        outputs, hidden = self.gru(d)

        return outputs, hidden


class BahdanauAttention(nn.Module):
    def __init__(self, hidden_size):
        super(BahdanauAttention, self).__init__()
        self.Wa = nn.Linear(hidden_size, hidden_size)
        self.Ua = nn.Linear(hidden_size, hidden_size)
        self.Va = nn.Linear(hidden_size, 1)

    def forward(self, query, keys):
        scores = self.Va(torch.tanh(self.Wa(query) + self.Ua(keys)))
        scores = scores.squeeze(2).unsqueeze(1)

        weights = F.softmax(scores, dim=-1)
        context = torch.bmm(weights, keys)

        return context, weights


class AttnDecoderRNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, dropout=0.1, max_length=10):
        super(AttnDecoderRNN, self).__init__()
        self.max_length = max_length
        self.attention = BahdanauAttention(hidden_size)
        self.gru = nn.GRU(input_size + hidden_size, hidden_size, batch_first=True)
        self.out = nn.Linear(hidden_size, output_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, encoder_outputs, encoder_hidden, embedded_target, embedded, is_train):
        decoder_input = embedded_target[:, 0].unsqueeze(1)

        decoder_hidden = encoder_hidden
        decoder_outputs = []
        attentions = []

        for i in range(self.max_length):
            decoder_output, decoder_hidden, attn_weights = self.forward_step(
                decoder_input, decoder_hidden, encoder_outputs
            )
            decoder_outputs.append(decoder_output)
            attentions.append(attn_weights)

            if is_train and i != self.max_length - 1:
                # Teacher forcing: Feed the target as the next input
                decoder_input = embedded_target[:, i + 1].unsqueeze(1)  # Teacher forcing
            else:
                # Without teacher forcing: use its own predictions as the next input
                _, topi = decoder_output.topk(1)
                decoder_input_i = topi.squeeze(-1).detach()  # detach from history as input
                decoder_input = embedded[decoder_input_i]

        decoder_outputs = torch.cat(decoder_outputs, dim=1)
        decoder_outputs = F.log_softmax(decoder_outputs, dim=-1)
        attentions = torch.cat(attentions, dim=1)

        return decoder_outputs, decoder_hidden, attentions

    def forward_step(self, input, hidden, encoder_outputs):
        dropout = self.dropout(input)

        query = hidden.permute(1, 0, 2)
        context, attn_weights = self.attention(query, encoder_outputs)
        input_gru = torch.cat((dropout, context), dim=2)

        output, hidden = self.gru(input_gru, hidden)
        output = self.out(output)

        return output, hidden, attn_weights


class GRU2GRU(nn.Module):
    def __init__(self, entity_dim, en_hidden_dim, de_hidden_dim,
                 de_output_dim, max_length=10, encoder_bidirectional=False):
        super(GRU2GRU, self).__init__()

        self.encoder = EncoderRNN(entity_dim, en_hidden_dim, bidirectional=encoder_bidirectional)
        self.decoder = AttnDecoderRNN(entity_dim, de_hidden_dim, de_output_dim, max_length=max_length)

    def forward(self, input_embedded, tgt_embedded, embedded, is_train):
        encoder_outputs, encoder_hidden = self.encoder(input_embedded)
        decoder_outputs, decoder_hidden, attentions = self.decoder(encoder_outputs, encoder_hidden, tgt_embedded, embedded,
                                                                   is_train=is_train)

        return decoder_outputs, decoder_hidden, attentions
